class Config():
    def __init__(self, mode):
        if mode == 'train':
            self._set_train_attr()
        elif mode == 'test':
            self._set_test_attr()

    def _set_train_attr(self):
        self.mode = 'train'
        self.no_accumulation = False
        self.data_name = 'ucf101'
        self.lstm_lr = 0.001
        self.hidden_size = 1024
        self.version = 'coviar-lstm'
        self.num_segments = 3
        self.representation = 'iframe'
        self.arch = 'resnet152'
        self.weights = None
        self.num_layers = 1
        self.dropout = 0
        self.data_root = '/home/zgp/Project/pytorch-coviar/data/ucf101/mpeg4_videos'
        self.train_list = "/home/zgp/Project/pytorch-coviar/data/datalists/ucf101_toy.txt"
        self.test_list = "/home/zgp/Project/pytorch-coviar/data/datalists/ucf101_toy.txt"
        self.batch_size = 60
        self.gpu = "0,1,2"
        self.workers = 8
        self.freeze = False
        self.lr = 0.001
        self.weight_decay = 0.1
        self.epochs = 10
        self.lr_steps = [150, 270, 390]
        self.lr_decay = 0.1
        self.eval_freq = 2
        # 智能生成
        self.logf = '/home/zgp/Project/coviar/out/dev/log/train.log'
        self.num_class = 101
        self.model_prefix = 'out/dev/weights/prefix'

    def _set_test_attr(self):
        self.mode = 'test'
        # disable accumulation of motion vectors and residuals.
        self.gpu: list = [0, 1, 2, 3]
        # number of workers for data loader.
        self.workers = 1
        self.no_accumulation: bool = False
        self.num_class = 101
        self.num_segments: int = 25
        self.num_layers = 1
        self.dropout = 0
        # Collect the feature input to LSTM and save to the dir lstm_infeature.
        self.lstm_infeature: str = '/home/zgp/Project/coviar/out/dev/lstmfeature'
        self.lstm_outfeature: str = '/home/zgp/Project/coviar/out/dev/lstmfeature'
        self.data_name = 'ucf101'
        self.hidden_size = 1024
        self.version = 'coviar-lstm'
        self.test_crops: int = 10
        self.batch_size = 1
        self.representation = 'iframe'
        self.arch = 'resnet152'
        self.weights = '/home/zgp/Project/coviar/out/coviar_lstm_ucf101_1/weights/ucf101_iframe_model_best.pth.tar'
        self.data_root = '/home/zgp/Project/pytorch-coviar/data/ucf101/mpeg4_videos'
        self.test_list = "/home/zgp/Project/pytorch-coviar/data/datalists/ucf101_toy.txt"
        self.save_scores: str = '/home/zgp/Project/coviar/out/dev/model_scores'
        self.logf = '/home/zgp/Project/coviar/out/dev/log/testlog.out'

    def __str__(self):
        info: str = ''
        for k, v in vars(self).items():
            info += f'\t{k}: {v}\n'

        return info
